/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \file 

    Definitions and explicit instatiations of the aux_grid class class
    used for making images of auxiliary quantities. */

// $Id$

#include <CCfits/CCfits>
#include "aux_grid.h"
#include "grid.h"
#include "arepo_grid.h"
#include "arepo_emission.h"

#ifdef WITH_AREPO
extern "C" {
#include "allvars.h"
}
#endif

using namespace std;

/** Constructor loads cell data from a FITS file.  (The grid structure
    is part of the base grid object and has already been loaded at
    this point.) */
template <template<typename> class grid_type,
	  template<typename> class sampling_policy,
	  typename rng_policy> 
mcrx::aux_grid<grid_type, sampling_policy, rng_policy>::
aux_grid (boost::shared_ptr<T_grid_impl> g,
	  CCfits::ExtHDU& structure_hdu,
	  CCfits::ExtHDU& data_hdu) :
  g_(g),
  emission_collection<generic_chromatic_policy<aux_pars_type>, 
  sampling_policy, rng_policy>()
{
  using namespace CCfits;
  
  // The grid structure already has to have been loaded upon
  // constructor of the grid_base object, so it's reasonable to expect
  // us to already have an open FITS file at this point.
  structure_hdu.readKey("lengthunit", units_ ["length"]);

  Column& c_m_g = data_hdu.column("mass_gas");
  Column& c_m_m = data_hdu.column("mass_metals");
  Column& c_sfr = data_hdu.column("SFR");
  Column& c_energy  = data_hdu.column("gas_temp_m");

  const int n = n_cells ();
  if (c_m_m.rows() != n ||
      c_m_g.rows() != n) {
    cerr << "Error: HDU does not contain the expected number of rows:\n "
	 << n << " cells in grid, but " << c_m_m.rows() 
	 << " rows in aux data"<< endl;
    return;
  }
    
  // read mass unit (length unit comes from the structure hdu).
  units_ ["mass"] = c_m_g.unit();

  std::vector<T_float> mass_gas;
  std::vector<T_float> mass_metals;
  std::vector<T_float> sfr;
  std::vector<T_float> energy;

  std::vector<T_float> L_bol;
  std::vector<T_float> mass_stars;
  std::vector<T_float> m_m_s;
  std::vector<T_float> age_m;
  std::vector<T_float> age_l;

  c_m_g.read(mass_gas, 1, n);
  c_m_m.read(mass_metals, 1, n);
  c_sfr.read(sfr, 1, n);
  c_energy.read(energy, 1, n);

  // If the emission data is in the grid, we use it.  Otherwise, we
  // put 0 in the emission part.
  try{
    Column& c_l_bol = data_hdu.column("L_bol");
    Column& c_ms = data_hdu.column("mass_stars");
    Column& c_mms = data_hdu.column("mass_stellar_metals");
    Column& c_agem = data_hdu.column("age_m");
    Column& c_agel = data_hdu.column("age_l");
    units_ ["luminosity"] = c_l_bol.unit();
    
    c_l_bol.read(L_bol, 1, n);
    c_ms.read(mass_stars, 1, n);
    c_mms.read(m_m_s, 1, n);
    c_agem.read(age_m, 1, n);
    c_agel.read(age_l, 1, n);
  }
  catch(Table::NoSuchColumn&) {
    L_bol.assign(n,0);
    mass_stars.assign(n,0);
    m_m_s.assign(n,0);
    age_m.assign(n,0);
    age_l.assign(n,0);
  }

  // This vector is used to collect the pointers to the emitters for
  // the call to distr_.setup.
  std::vector<emission<generic_chromatic_policy<T_content>,
    rng_policy>*> pointers;
  array_1 weights(n);
  
  int i = 0;
  for (iterator c = begin (); 
       c != end(); ++c, ++i) {
    // Note the normalization, the emitter contents are not the
    // quantities themselves but quantity/emission weight.  The
    // emission weight of the cells is stellar and gas mass, and we need to
    // take that out.  Final normalization factor is
    // sum(emission weight)/N_rays, which needs to be applied to the
    // cameras afterwards.

    // these are decoded in aux_emergence::save_images.
    T_content d;
    d[aux_pars_fields::mass_gas] = mass_gas [i];
    d[aux_pars_fields::mass_metals] = mass_metals [i];
    d[aux_pars_fields::SFR] = sfr [i];
    d[aux_pars_fields::energy] = energy [i];
    d[aux_pars_fields::mass_stars] = mass_stars [i];
    d[aux_pars_fields::mass_metals_stars] = m_m_s [i];
    d[aux_pars_fields::L_bol] = L_bol [i];
    d[aux_pars_fields::age_m] = (age_m [i]==age_m [i] ? age_m[i] : 0);
    d[aux_pars_fields::age_l] = (age_l [i]==age_l [i] ? age_l[i] : 0);
    d[aux_pars_fields::FIR] = 0; // done in separate stage
   
    const T_float emission_weight = mass_gas [i] + mass_stars [i];
    if(emission_weight!=0)
      d*= 1/emission_weight;
    else
      // to avoid NaN, we must set emission to zero in this case
      d=0;

    typename T_grid::T_data* dd = new T_emitter (d, *c);
    c->set_data(dd);

    // save variables for setup of emission_collection 
    pointers.push_back(dd);
    weights(i) = emission_weight;
  }

  // init distribution
  this->distr_.setup(pointers, weights);
}


/** This constructor loads cell data from the Arepo data
    structures. The grid structure is handled by the supplied grid
    object, all that we do here is use the SPH data to populate our
    data. */
template <template<typename> class grid_type,
	  template<typename> class sampling_policy,
	  typename rng_policy> 
mcrx::aux_grid<grid_type, sampling_policy, rng_policy>::
aux_grid (boost::shared_ptr<T_grid_impl> g) :
  g_(g),
  emission_collection<generic_chromatic_policy<aux_pars_type>, 
  sampling_policy, rng_policy>()
{
#ifdef WITH_AREPO
  const int nc = this->n_cells ();
  assert(nc==N_gas);
  cout << "Creating cell data from Arepo data for " << nc << " cells. " << endl;

  // First copy units from the arepo units
  units_["length"] = arepo::arepo_units.find("length")->second;
  units_["time"] = arepo::arepo_units.find("time")->second;
  units_["mass"] = arepo::arepo_units.find("mass")->second;

  // to convert density to our units
  const T_float convert_density =
    All.UnitDensity_in_cgs*units::convert("g/cm^3",units_["mass"]+"/"+units_["length"]+"^3");
  const T_float convert_sfr =
    All.UnitMass_in_g/All.UnitTime_in_s*
    units::convert("g/s",units_["mass"]+"/"+units_["time"]);

  // This vector is used to collect the pointers to the emitters for
  // the call to distr_.setup.
  std::vector<emission<generic_chromatic_policy<T_content>,
    rng_policy>*> pointers;
  array_1 weights(nc);
  
  int i = 0;
  for (iterator c = begin (); c != end(); ++c, ++i) {
    // Note the normalization, the emitter contents are not the
    // quantities themselves but quantity/emission weight.  The
    // emission weight of the cells is stellar and gas mass, and we need to
    // take that out.  Final normalization factor is
    // sum(emission weight)/N_rays, which needs to be applied to the
    // cameras afterwards.

    // these are decoded in aux_emergence::save_images.
    T_content d;
    d[aux_pars_fields::mass_gas] = SphP [i].Density*c->volume()*convert_density;
    d[aux_pars_fields::mass_metals] = d[aux_pars_fields::mass_gas]*P[i].Metallicity;
    d[aux_pars_fields::SFR] = SphP[i].Sfr*convert_sfr; //how do we get sfr?

    // use Arepo to calculate temperature directly (hopefully it outputs in K...)
    T_float ne=SphP[i].Ne;
    const double T= convert_u_to_temp(SphP[i].Utherm, SphP [i].Density, &ne);
    d[aux_pars_fields::energy] = T*d[aux_pars_fields::mass_gas];

    // all of these are from particles, not the grid.
    d[aux_pars_fields::mass_stars] = 0;
    d[aux_pars_fields::mass_metals_stars] = 0;
    d[aux_pars_fields::L_bol] = 0;
    d[aux_pars_fields::age_m] = 0;
    d[aux_pars_fields::age_l] = 0;
    d[aux_pars_fields::FIR] = 0; 
   
    const T_float emission_weight = 
      d[aux_pars_fields::mass_gas] + d[aux_pars_fields::mass_stars];
    if(emission_weight!=0)
      d*= 1/emission_weight;
    else
      // to avoid NaN, we must set emission to zero in this case
      d=0;

    typename T_grid::T_data* dd = new T_emitter (d, *c);
    c->set_data(dd);

    // save variables for setup of emission_collection 
    pointers.push_back(dd);
    weights(i) = emission_weight;
  }

  // init distribution
  this->distr_.setup(pointers, weights);

  // this call goes to the emission_collection class and makes it
  // possible to do intelligent stuff after the individual emitters
  // have been set up
  T_emitter::post_setup(*g_);

#else
  cerr << "Not compiled with Arepo support" << endl;
  exit(1);
#endif
}

namespace mcrx {
  template class aux_grid<adaptive_grid, cumulative_sampling, local_random>;
#ifdef WITH_AREPO
  template class aux_grid<arepo_grid, cumulative_sampling, local_random>;
#endif
}

